from sqlalchemy import (
    Column,
    UUID,
    text,
    String,
    DateTime,
    Boolean,
    Integer,
    Text, PrimaryKeyConstraint, func
)
from sqlalchemy.dialects.postgresql import JSONB

from internal.extension import db


class DataSet(db.Model):
    """知识库表"""
    __tablename__ = 'dataset'
    __table_args__ = (
        PrimaryKeyConstraint('id', name='dataset_id_pkey'),
    )
    id = Column(UUID, nullable=False, server_default=text('uuid_generate_v4()'))
    name = Column(String, nullable=False, server_default=text("''::character varying"))
    icon = Column(String, nullable=False, server_default=text("''::character varying"))
    description = Column(Text, nullable=True, server_default=text("''::text"))
    updated_at = Column(
        DateTime,
        nullable=False,
        server_default=text('CURRENT_TIMESTAMP(0)'),
        server_onupdate=text('CURRENT_TIMESTAMP(0)'))
    created_at = Column(
        DateTime,
        nullable=False,
        server_default=text('CURRENT_TIMESTAMP(0)')
    )

    @property
    def get_character_count(self):
        return (db.session.query(func.coalesce(func.sum(Document.character_count), 0)).
                filter(Document.dataset_id == self.id)).scalar()

    @property
    def get_document_count(self):
        return (db.session.query(func.count(Document.id)).
                filter(Document.dataset_id == self.id)).scalar()

    @property
    def get_hit_count(self):
        return (db.session.query(func.coalesce(func.sum(Segment.hit_count), 0)).
                filter(Segment.dataset_id == self.id)).scalar()

    @property
    def get_app_count(self):
        return (db.session.query(func.count(AppDatasetJoin.id)).
                filter(AppDatasetJoin.dataset_id == self.id)).scalar()


class Document(db.Model):
    """文档表模型"""
    __tablename__ = 'document'
    __table_args__ = (
        PrimaryKeyConstraint('id', name='document_id_pkey'),
    )
    id = Column(UUID, nullable=False, server_default=text('uuid_generate_v4()'))
    account_id = Column(UUID, nullable=False)
    dataset_id = Column(UUID, nullable=False)
    upload_file_id = Column(UUID, nullable=False)
    process_rule_id = Column(UUID, nullable=False)
    bash = Column(String, nullable=False, server_default=text("''::character varying"))
    name = Column(String, nullable=False, server_default=text("''::character varying"))
    position = Column(Integer, nullable=False, server_default=text("1"))
    character_count = Column(Integer, nullable=False, server_default=text("0"))
    token_count = Column(Integer, nullable=False, server_default=text("0"))
    processing_started_at = Column(DateTime, nullable=True)  # 处理时间
    parseing_completed_at = Column(DateTime, nullable=True)  # 解析时间
    splitting_completed_at = Column(DateTime, nullable=True)  # 文档切割时间
    indexing_completed_at = Column(DateTime, nullable=True)  # 构建索引时间
    completed_at = Column(DateTime, nullable=True)  # 完成时间
    stopped_at = Column(DateTime, nullable=True)  # 停止时间
    error = Column(Text, nullable=True, server_default=text("''::text"))  # 发生错误的内容
    enabled = Column(Boolean, nullable=False, server_default=text("false"))  # 是否启用
    disabled_at = Column(DateTime, nullable=True)  # 禁用时间
    status = Column(String(255), nullable=False, server_default=text("'waiting'::character varying"))  # 状态
    updated_at = Column(
        DateTime,
        nullable=False,
        server_default=text('CURRENT_TIMESTAMP(0)'),
        server_onupdate=text('CURRENT_TIMESTAMP(0)'))
    created_at = Column(
        DateTime,
        nullable=False,
        server_default=text('CURRENT_TIMESTAMP(0)')
    )


class Segment(db.Model):
    """文档分片表模型"""
    __tablename__ = 'segment'
    __table_args__ = (
        PrimaryKeyConstraint('id', name='segment_id_pkey'),
    )
    id = Column(UUID, nullable=False, server_default=text('uuid_generate_v4()'))
    account_id = Column(UUID, nullable=False)
    dataset_id = Column(UUID, nullable=False)
    document_id = Column(UUID, nullable=False)
    node_id = Column(UUID, nullable=False)
    position = Column(Integer, nullable=False, server_default=text("1"))
    content = error = Column(Text, nullable=True, server_default=text("''::text"))  # 发生错误的内容
    character_count = Column(Integer, nullable=False, server_default=text("0"))
    token_count = Column(Integer, nullable=False, server_default=text("0"))
    keywords = Column(JSONB, nullable=False, server_default=text("'[]'::jsonb"))
    hash = Column(String(255), nullable=False, server_default=text("''::character varying"))
    hit_count = Column(Integer, nullable=False, server_default=text("0"))
    enabled = Column(Boolean, nullable=False, server_default=text("false"))
    disabled_at = Column(DateTime, nullable=True)
    processing_started_at = Column(DateTime, nullable=True)
    indexing_completed_at = Column(DateTime, nullable=True)
    completed_at = Column(DateTime, nullable=True)
    status = Column(String(255), nullable=False, server_default=text("'waiting'::character varying"))  # 状态
    updated_at = Column(
        DateTime,
        nullable=False,
        server_default=text('CURRENT_TIMESTAMP(0)'),
        server_onupdate=text('CURRENT_TIMESTAMP(0)'))
    created_at = Column(
        DateTime,
        nullable=False,
        server_default=text('CURRENT_TIMESTAMP(0)')
    )


class KeywordTable(db.Model):
    """关键词表模型"""
    __tablename__ = 'keyword_table'
    __table_args__ = (
        PrimaryKeyConstraint('id', name='keyword_id_pkey'),
    )
    id = Column(UUID, nullable=False, server_default=text('uuid_generate_v4()'))
    account_id = Column(UUID, nullable=False)
    dataset_id = Column(UUID, nullable=False)
    keyword_table = Column(JSONB, nullable=False, server_default=text("'{}'::jsonb"))
    updated_at = Column(
        DateTime,
        nullable=False,
        server_default=text('CURRENT_TIMESTAMP(0)'),
        server_onupdate=text('CURRENT_TIMESTAMP(0)'))
    created_at = Column(
        DateTime,
        nullable=False,
        server_default=text('CURRENT_TIMESTAMP(0)')
    )


class DatasetQuery(db.Model):
    """知识库查询表"""
    __tablename__ = 'dataset_query'
    __table_args__ = (
        PrimaryKeyConstraint('id', name='dataset_query_id_pkey'),
    )
    id = Column(UUID, nullable=False, server_default=text('uuid_generate_v4()'))
    account_id = Column(UUID, nullable=False)
    dataset_id = Column(UUID, nullable=False)
    query = Column(Text, nullable=False, server_default=text("''::text"))
    source = Column(String(255), nullable=False, server_default=text("'HitTesting'::character varying"))  # 默认来源为召回测试
    source_app = Column(UUID, nullable=False)
    created_by = Column(UUID, nullable=False)
    updated_at = Column(
        DateTime,
        nullable=False,
        server_default=text('CURRENT_TIMESTAMP(0)'),
        server_onupdate=text('CURRENT_TIMESTAMP(0)'))
    created_at = Column(
        DateTime,
        nullable=False,
        server_default=text('CURRENT_TIMESTAMP(0)')
    )


class ProcessRule(db.Model):
    """知识库规则表"""
    __tablename__ = 'process_rule'
    __table_args__ = (
        PrimaryKeyConstraint('id', name='process_rule_id_pkey'),
    )
    id = Column(UUID, nullable=False, server_default=text('uuid_generate_v4()'))
    account_id = Column(UUID, nullable=False)
    dataset_id = Column(UUID, nullable=False)
    mode = Column(String(255), nullable=False, server_default=text("'automic'::character varying"))
    rule = Column(JSONB, nullable=False, server_default=text("'{}'::jsonb"))
    updated_at = Column(
        DateTime,
        nullable=False,
        server_default=text('CURRENT_TIMESTAMP(0)'),
        server_onupdate=text('CURRENT_TIMESTAMP(0)'))
    created_at = Column(
        DateTime,
        nullable=False,
        server_default=text('CURRENT_TIMESTAMP(0)')
    )


class AppDatasetJoin(db.Model):
    """应用与知识库关联表"""
    __tablename__ = 'app_dataset_join'
    __table_args__ = (
        PrimaryKeyConstraint('id', name='app_dataset_join_id_pkey'),
    )
    id = Column(UUID, nullable=False, server_default=text('uuid_generate_v4()'))
    app_id = Column(UUID, nullable=False)
    dataset_id = Column(UUID, nullable=False)
    updated_at = Column(
        DateTime,
        nullable=False,
        server_default=text('CURRENT_TIMESTAMP(0)'),
        server_onupdate=text('CURRENT_TIMESTAMP(0)'))
    created_at = Column(
        DateTime,
        nullable=False,
        server_default=text('CURRENT_TIMESTAMP(0)')
    )
